"""
Kraken Flux Empty Latent Image
MIT License © 2025 The Kraken (@KrakenUnbound)
This node produces a Flux-style empty latent with a reliable WxH preview and
emits both a normalized `preset_out` string ("WxH") and a human `resolution` ("W x H").
Designed for clean hand-offs into Upscale & Tile Calc and Resolution Helper.
"""

import torch
import math
import os
import time
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import folder_paths

counter = 0  # Global counter for unique filenames

class KrakenFluxEmptyLatentImage:
    @classmethod
    def INPUT_TYPES(cls):
        megapixel_options = ["0.5", "0.75", "1.0", "1.25", "1.5", "1.75", "2.0", "2.25", "2.5", "2.75", "3.0", "Custom"]
        # Combined aspect_ratio presets with use-case resolutions
        aspect_options = [
            "Custom",
            "1:1 (Perfect Square)",
            "2:1 (Panorama Landscape)",
            "2:3 (Classic Portrait)",
            "3:4 (Golden Ratio)",
            "3:5 (Elegant Vertical)",
            "4:5 (Artistic Frame)",
            "5:7 (Balanced Portrait)",
            "5:8 (Tall Portrait)",
            "7:9 (Modern Portrait)",
            "9:16 (Slim Vertical)",
            "9:19 (Tall Slim)",
            "9:21 (Ultra Tall)",
            "9:32 (Skyline)",
            "3:2 (Golden Landscape)",
            "4:3 (Classic Landscape)",
            "5:3 (Wide Horizon)",
            "5:4 (Balanced Frame)",
            "7:5 (Elegant Landscape)",
            "8:5 (Cinematic View)",
            "9:7 (Artful Horizon)",
            "16:9 (Panorama)",
            "19:9 (Cinematic Ultrawide)",
            "21:9 (Epic Ultrawide)",
            "32:9 (Extreme Ultrawide)",
            "Screen: iPhone Wallpaper - 1170x2532 (9:16)",
            "Screen: iPhone Pro Wallpaper - 1290x2796 (9:19.5)",
            "Screen: Android HD Wallpaper - 1080x1920 (9:16)",
            "Screen: Android QHD Wallpaper - 1440x3200 (9:20)",
            "Screen: Desktop HD Wallpaper - 1920x1080 (16:9)",
            "Screen: Desktop QHD Wallpaper - 2560x1440 (16:9)",
            "Screen: Dual Monitor Wallpaper - 5120x1440 (32:9)",
            "Screen: 4K Monitor - 3840x2160 (16:9)",
            "Screen: Ultrawide QHD - 3440x1440 (21:9)",
            "Screen: Ultrawide 4K - 3840x1600 (12:5)",
            "Video: KlingAI Landscape - 1920x1080 (16:9)",
            "Video: KlingAI Portrait - 1080x1920 (9:16)",
            "Video: KlingAI Square - 1024x1024 (1:1)",
            "Video: KlingAI Wide - 1344x768 (3:2)",
            "Video: KlingAI Tall - 768x1344 (2:3)",
            "Social: Instagram Post Square - 1080x1080 (1:1)",
            "Social: Instagram Post Portrait - 1080x1350 (4:5)",
            "Social: Instagram Story - 1080x1920 (9:16)",
            "Social: YouTube Thumbnail - 1280x720 (16:9)",
            "Social: Facebook Post - 1200x1200 (1:1)",
            "Social: Facebook Banner - 1200x628 (1.91:1)",
            "Social: Twitter/X Post - 1200x675 (16:9)",
            "Social: LinkedIn Post - 1200x1200 (1:1)",
            "Advertising: Leaderboard Banner - 728x90 (8:1)",
            "Advertising: Medium Rectangle Banner - 300x250 (1.2:1)",
            "Advertising: Large Rectangle Banner - 336x280 (1.2:1)",
            "Advertising: Full Banner - 468x60 (7.8:1)",
            "Advertising: Wide Skyscraper Banner - 160x600 (1:3.75)",
            "Advertising: Half Page Banner - 300x600 (1:2)",
            "Advertising: Large Leaderboard Banner - 970x90 (10.8:1)",
            "Advertising: Square Banner - 250x250 (1:1)",
            "Print: Movie Poster - 8100x12000 (27x40in at 300 DPI)",
            "Print: Flyer Letter - 2550x3300 (8.5x11in at 300 DPI)",
            "Print: Flyer Legal - 2550x4200 (8.5x14in at 300 DPI)",
            "Print: Business Card - 1050x600 (3.5x2in at 300 DPI)",
            "Print: Poster A3 - 3508x4961 (11.7x16.5in at 300 DPI)",
            "Print: Poster A2 - 4961x7016 (16.5x23.4in at 300 DPI)",
            "Print: Poster A1 - 7016x9933 (23.4x33.1in at 300 DPI)",
            "Print: Poster A0 - 9933x14043 (33.1x46.8in at 300 DPI)"
        ]
        divisible_options = ["Auto", 8, 16, 32, 64, 128, 256]
        return {
            "required": {
                "megapixel": (megapixel_options, {"default": "1.0"}),
                "preset": (aspect_options, {"default": "Custom"}),
                "use_custom_aspect": ("BOOLEAN", {"default": False}),
                "custom_aspect": ("STRING", {"default": "1:1"}),
                "width": ("INT", {"default": 1024, "min": 64, "max": 4096, "step": 8}),
                "height": ("INT", {"default": 1024, "min": 64, "max": 4096, "step": 8}),
                "divisible_by": (divisible_options, {"default": "Auto"}),
                "batch_size": ("INT", {"default": 1, "min": 1, "max": 64}),
            }
        }

    RETURN_TYPES = ("LATENT", "INT", "INT", "STRING", "STRING")
    RETURN_NAMES = ("latent", "width_out", "height_out", "resolution", "preset_out")
    FUNCTION = "generate"
    CATEGORY = "kraken/latent"

    
    
    
    
    
    def create_preview_image(self, width, height, goal_str, ratio_display):
        # Preview rendered at EXACT latent size so ComfyUI shows correct WxH at the bottom.
        from PIL import Image, ImageDraw, ImageFont

        def load_font(sz):
            for name in ("DejaVuSans.ttf", "Arial.ttf", "arial.ttf"):
                try:
                    return ImageFont.truetype(name, sz)
                except Exception:
                    continue
            return ImageFont.load_default()

        def text_wh(draw, text, font):
            try:
                l, t, r, b = draw.textbbox((0, 0), text, font=font)
                return (r - l, b - t)
            except Exception:
                try:
                    return font.getsize(text)
                except Exception:
                    sz = getattr(font, "size", 18)
                    return (int(len(text) * sz * 0.6), int(sz * 1.2))

        w, h = int(width), int(height)
        image = Image.new('RGB', (w, h), (24, 24, 24))
        draw  = ImageDraw.Draw(image)

        # Border
        draw.rectangle([0, 0, w - 1, h - 1], outline=(160, 160, 160), width=2)

        # Font sizes relative to canvas
        base = max(20, int(min(w, h) * 0.14))   # big center size text
        sub  = max(14, int(base * 0.60))        # info text
        font_main = load_font(base)
        font_sub  = load_font(sub)

        # Big centered latent size
        center_text = f"{w}x{h}"
        tw, th = text_wh(draw, center_text, font_main)
        cx, cy = w // 2, h // 2
        # outline for readability
        for off in ((-2,-2),(2,2),(-2,2),(2,-2)):
            draw.text((cx - tw//2 + off[0], cy - th//2 + off[1]), center_text, fill=(0,0,0), font=font_main)
        draw.text((cx - tw//2, cy - th//2), center_text, fill=(245,245,245), font=font_main)

        # Top-left info block: show target if present; otherwise show ratio
        lines = []
        if goal_str:
            lines.append(f"target {goal_str}")   # e.g., 1080x1920
        # Show ratio if it looks like A:B or matches known form
        if isinstance(ratio_display, str) and (':' in ratio_display or 'x' in ratio_display):
            lines.append(f"ratio  {ratio_display}")

        y = 10  # add more padding so it doesn't collide with border
        for line in lines:
            lw, lh = text_wh(draw, line, font_sub)
            pad = 5
            # dark backdrop
            draw.rectangle([10 - pad, y - pad, 10 + lw + pad, y + lh + pad], fill=(0,0,0))
            # draw text
            draw.text((10, y), line, fill=(220,220,220), font=font_sub)
            y += lh + 8

        return image


    def generate(self, megapixel, preset, use_custom_aspect, custom_aspect, width, height, divisible_by, batch_size):
        ar = 1.0
        ratio_display = f"{width}x{height}"
        if use_custom_aspect:
            ratio_display = custom_aspect
            if ':' in custom_aspect:
                w_part, h_part = map(float, custom_aspect.split(":"))
                ar = w_part / h_part
            else:
                ar = float(custom_aspect)
        elif preset != "Custom":
            if " - " in preset:
                # Specific resolution preset
                size_str = preset.split(" - ")[1].split(" (")[0]
                width, height = map(int, size_str.split("x"))
                ar = width / height
                ratio_display = size_str  # Use just WIDTHxHEIGHT for display
            else:
                # Ratio-based preset
                ratio_str = preset.split(" (")[0]
                w_part, h_part = map(float, ratio_str.split(":"))
                ar = w_part / h_part
                ratio_display = f"{w_part}:{h_part}"
        else:
            ar = width / height
            ratio_display = f"{width}x{height}"

        if megapixel == "Custom":
            mp = (width * height) / 1_000_000
            use_calculated = False
        else:
            mp = float(megapixel)
            use_calculated = True

        if use_calculated:
            base_width = math.sqrt(mp * 1_000_000 * ar)
            base_height = math.sqrt(mp * 1_000_000 / ar)
        else:
            base_width = width
            base_height = height

        if divisible_by == "Auto":
            if mp > 1.5:
                divisible_by = 64
            elif mp > 0.75:
                divisible_by = 32
            else:
                divisible_by = 16
        else:
            divisible_by = int(divisible_by)

        width = round(base_width / divisible_by) * divisible_by
        height = round(base_height / divisible_by) * divisible_by

        if width % 8 != 0 or height % 8 != 0:
            width = (width // 8) * 8
            height = (height // 8) * 8

        latent = torch.zeros([batch_size, 16, height // 8, width // 8])

        resolution = f"{width} x {height}"
        # Determine preset_out: prefer explicit target size from preset label if present
        target_preset_str = None
        try:
            if preset != "Custom" and " - " in preset:
                target_preset_str = preset.split(" - ")[1].split(" (")[0]
        except Exception:
            target_preset_str = None
        preset_out = target_preset_str if target_preset_str else f"{width}x{height}"  # WIDTHxHEIGHT string for downstream
        # Determine preset_out: prefer explicit target size from preset label if present

        preview_img = self.create_preview_image(width, height, target_preset_str or '', ratio_display)

        global counter
        counter += 1
        full_output_folder = folder_paths.get_temp_directory()
        current_time = time.time()
        file = f"kraken_preview_{current_time}_{counter:05}.png"
        file_path = os.path.join(full_output_folder, file)
        preview_img.save(file_path, compress_level=4)

        images = [{"filename": file, "subfolder": "", "type": "temp"}]

        return {
            "ui": {"images": images},
            "result": ({"samples": latent}, width, height, resolution, preset_out)
        }